# !/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : 批量写入数据(helpers.async_bulk).py
# @Author: dongguangwen
# @Date  : 2025-05-31 19:00
import asyncio
from elasticsearch import AsyncElasticsearch
from elasticsearch.helpers import async_bulk
import json
import time

# 创建异步客户端连接
es_async = AsyncElasticsearch(hosts=['http://192.168.1.119:9200'])


async def load_data_and_bulk():
    actions = []
    with open('./person_info_100w.json', 'r', encoding='utf-8') as f:
        for line in f:
            data = json.loads(line.strip())
            action = {
                "_index": "persons",
                "_id": data.get('id'),
                "_source": {
                    "id": data.get('id'),
                    "name": data.get('name'),
                    "sex": data.get('sex'),
                    "age": data.get('age'),
                    "email": data.get('email'),
                    "simple_profile": {
                        "username": data['simple_profile'].get('username'),
                        "name": data['simple_profile'].get('name'),
                        "sex": data['simple_profile'].get('sex'),
                        "address": data['simple_profile'].get('address'),
                        "mail": data['simple_profile'].get('mail'),
                        "birthdate": data['simple_profile'].get('birthdate'),
                    },
                    "profile": {
                        "job": data['profile'].get('job'),
                        "company": data['profile'].get('company'),
                        "ssn": data['profile'].get('ssn'),
                        "residence": data['profile'].get('residence'),
                        "blood_group": data['profile'].get('blood_group'),
                        "website": data['profile'].get('website'),
                        "username": data['profile'].get('username'),
                        "name": data['profile'].get('name'),
                        "sex": data['profile'].get('sex'),
                        "address": data['profile'].get('address'),
                        "mail": data['profile'].get('mail'),
                        "birthdate": data['profile'].get('birthdate'),
                    },
                    "location": {
                        "country": data.get('country'),
                        "province": data.get('province'),
                        "city": data.get('city'),
                        "street_name": data.get('street_name'),
                        "street_address": data.get('street_address'),
                        "address": data.get('address'),
                        "postcode": data.get('postcode'),
                        "gps": data.get('gps')
                    },
                    "net_info": {
                        "ipv4": data.get('ipv4'),
                        "ipv6": data.get('ipv6'),
                        "uri": data.get('uri'),
                        "url": data.get('url'),
                        "img_url": data.get('img_url'),
                        "domain": data.get('domain'),
                        "user_agent": data.get('user_agent')
                    },
                    "character": data.get('character'),
                    "long_text": data.get('long_text'),
                    "time": data.get('time')
                }
            }
            actions.append(action)

            # 每批发送 300 条
            if len(actions) >= 300:
                await async_bulk(es_async, actions)
                print(f"已写入 {len(actions)} 条数据")
                actions = []

        # 处理剩余数据
        if actions:
            await async_bulk(es_async, actions)
            print(f"最后一批写入 {len(actions)} 条数据")


async def main():
    start = time.time()
    await load_data_and_bulk()
    end = time.time()
    print("总耗时：", end - start, "秒")

if __name__ == "__main__":
    asyncio.run(main())

"""
总耗时： 593.7879345417023 秒
"""